import keras
from bilstm_crf_model import BilstmCrfModel
from data_processing import DataProcessor
from metrics import *
import pickle
from sklearn.model_selection import train_test_split
max_len = 80
vocab_size = 6000
embedding_dim = 300
lstm_units = 128

if __name__ == '__main__':
    data = DataProcessor(max_len,vocab_size)

    train_x, trainy,test_x, testy = data.read_data()

    train_x, trainy = data.encoder(train_x,trainy)

    x_train, x_dev, y_train, y_dev = train_test_split(train_x, trainy, test_size=0.25)



    test_x, testy = data.encoder(test_x,testy)

    class_nums = data.class_nums
    word2id = data.word2id
    tag2id = data.tag2id
    id2tag = data.id2tag

    pickle.dump(word2id,open("checkpoint/word2id.pkl","wb"))
    pickle.dump(tag2id,open("checkpoint/tag2id.pkl","wb"))
    pickle.dump(id2tag,open("checkpoint/id2tag.pkl","wb"))

    bilstm_crf = BilstmCrfModel(
        max_len=max_len,
        vocab_size=vocab_size,
        embedding_dim=embedding_dim,
        lstm_units=lstm_units,
        class_nums=class_nums
    )

    model= bilstm_crf.buid_model()

    reduce_lr = keras.callbacks.ReduceLROnPlateau(
        monitor='val_loss',
        factor=0.5,
        patience=4,
        verbose=1)

    earlystop = keras.callbacks.EarlyStopping(
        monitor='val_loss',
        patience=10,
        verbose=2,
        mode='min'
    )

    bast_model_filepath = 'checkpoint/best_bilstm_crf_model.h5'
    checkpoint = keras.callbacks.ModelCheckpoint(
        bast_model_filepath,
        monitor='val_loss',
        verbose=5,
        save_best_only=True,
        mode='min'
    )
    model.fit(
        x=x_train,
        y=y_train,
        batch_size=64,
        epochs=20,
        validation_data=(x_dev, y_dev),
        shuffle=True,
        callbacks=[reduce_lr, earlystop, checkpoint]
    )

    model.load_weights(bast_model_filepath)

    model.save('checkpoint/bilstm_crf_model.h5')

    pred = model.predict(test_x)

    y_true, y_pred = [], []



    for t_oh, p_oh in zip(testy, pred):
        t_oh = np.argmax(t_oh, axis=1)
        t_oh = [id2tag[i].replace('_', '-') for i in t_oh if i != 0]
        p_oh = np.argmax(p_oh, axis=1)
        p_oh = [id2tag[i].replace('_', '-') for i in p_oh if i != 0]

        y_true.append(t_oh)
        y_pred.append(p_oh)

    f1 = f1_score(y_true, y_pred, suffix=False)
    p = precision_score(y_true, y_pred, suffix=False)
    r = recall_score(y_true, y_pred, suffix=False)
    acc = accuracy_score(y_true, y_pred)
    print(
        "f1_score: {:.4f}, precision_score: {:.4f}, recall_score: {:.4f}, accuracy_score: {:.4f}".format(f1, p, r, acc))
    print(classification_report(y_true, y_pred, digits=4, suffix=False))
